# Custom sensitivity analysis table function
# Alex Gazmararian
# agazmararian@gmail.com

#' Generate customized sensitivity analysis table with caption at top and resized table
#'
#' @param sense_model A sensemakr object from the sensemakr() function
#' @param outcome_label String label for the outcome variable
#' @param treatment_label String label for the treatment variable  
#' @param file_name String path where to save the .tex file
#' @param caption String caption text for the table
#' @param label String LaTeX label for the table (without "tab:" prefix)
#' @param resize_width Numeric proportion of textwidth for table size (default 0.8)
#' @param digits Integer number of digits for rounding (default 2)
#' @param verbose Logical whether to show verbose output (default FALSE)
#'
#' @return Invisibly returns the modified LaTeX string
#' @export
#'
#' @examples
#' \dontrun{
#' # Generate sensitivity analysis
#' sense_out <- sensemakr(...)
#' 
#' # Create custom table
#' generate_sensitivity_table(
#'   sense_model = sense_out,
#'   outcome_label = "Recognition (=1)",
#'   treatment_label = "Manufacturing proximity (Q1)",
#'   file_name = here("output", "tables", "sensitivity_mfg.tex"),
#'   caption = "Sensitivity analysis for manufacturing proximity model",
#'   label = "sensitivity_mfg"
#' )
#' }
generate_sensitivity_table <- function(sense_model, 
                                     outcome_label, 
                                     treatment_label,
                                     file_name,
                                     caption,
                                     label,
                                     resize_width = 0.8,
                                     digits = 2,
                                     verbose = FALSE) {
  
  # Validate inputs
  if (!inherits(sense_model, "sensemakr")) {
    stop("sense_model must be a sensemakr object")
  }
  
  if (!is.character(outcome_label) || length(outcome_label) != 1) {
    stop("outcome_label must be a single character string")
  }
  
  if (!is.character(treatment_label) || length(treatment_label) != 1) {
    stop("treatment_label must be a single character string")
  }
  
  if (!is.character(file_name) || length(file_name) != 1) {
    stop("file_name must be a single character string")
  }
  
  if (!is.character(caption) || length(caption) != 1) {
    stop("caption must be a single character string")
  }
  
  if (!is.character(label) || length(label) != 1) {
    stop("label must be a single character string")
  }
  
  if (!is.numeric(resize_width) || resize_width <= 0 || resize_width > 1) {
    stop("resize_width must be a number between 0 and 1")
  }
  
  # Load required packages
  if (!requireNamespace("sensemakr", quietly = TRUE)) {
    stop("sensemakr package is required")
  }
  
  if (!requireNamespace("stringr", quietly = TRUE)) {
    stop("stringr package is required")
  }
  
  # Generate the basic table using ovb_minimal_reporting
  if (verbose) {
    message("Generating basic sensitivity table...")
  }
  
  tex_output <- sensemakr::ovb_minimal_reporting(
    sense_model,
    format = "latex",
    digits = digits, 
    outcome_label = outcome_label,
    treatment_label = treatment_label,
    verbose = verbose
  )
  
  if (verbose) {
    message("Modifying table structure...")
  }
  
  # Helper function to format a single number with consistent decimal places
  format_num <- function(x, d) {
    formatC(as.numeric(x), format = "f", digits = d)
  }
  
  # Post-process to ensure consistent digit formatting
  # 1. Handle percentages: find "N\%" or "N.N\%" and reformat to d decimal places
  tex_output <- gsub(
    "([0-9]+\\.?[0-9]*)\\\\%",
    "<<PCT:\\1>>",  # temporary placeholder to avoid double-processing
    tex_output
  )
  # Now find and replace the placeholders with properly formatted values
  while (grepl("<<PCT:", tex_output)) {
    m <- regexpr("<<PCT:([0-9]+\\.?[0-9]*)>>", tex_output, perl = TRUE)
    if (m == -1) break
    full_match <- regmatches(tex_output, m)
    num_val <- sub("<<PCT:(.*)>>", "\\1", full_match)
    formatted <- paste0(format_num(num_val, digits), "\\%")
    tex_output <- sub(full_match, formatted, tex_output, fixed = TRUE)
  }
  
  # 2. Handle regular decimal numbers in table cells (between & and & or & and \\)
  # Pattern: & followed by a decimal number followed by & or \\
  # Replace numbers with fewer than 'digits' decimal places
  tex_output <- gsub(
    "& ([0-9]+\\.[0-9]+) &",
    "<<NUM:\\1>>&",
    tex_output
  )
  tex_output <- gsub(
    "& ([0-9]+\\.[0-9]+) \\\\\\\\",
    "<<NUM:\\1>>\\\\\\\\",
    tex_output
  )
  # Process the number placeholders
  while (grepl("<<NUM:", tex_output)) {
    m <- regexpr("<<NUM:([0-9]+\\.[0-9]+)>>", tex_output, perl = TRUE)
    if (m == -1) break
    full_match <- regmatches(tex_output, m)
    num_val <- sub("<<NUM:(.*)>>", "\\1", full_match)
    formatted <- paste0("& ", format_num(num_val, digits), " ")
    tex_output <- sub(full_match, formatted, tex_output, fixed = TRUE)
  }
  
  # Parse and modify the LaTeX output
  tex_lines <- strsplit(tex_output, "\n")[[1]]
  new_lines <- character()
  
  # Process each line to customize the table
  for(i in 1:length(tex_lines)) {
    if(grepl("^\\\\begin\\{table\\}", tex_lines[i])) {
      # Replace table beginning with custom formatting
      new_lines <- c(new_lines, "\\begin{table}[!h]")
    } else if(grepl("^\\\\centering", tex_lines[i])) {
      # Add caption, label, centering, and resizebox
      new_lines <- c(new_lines, 
                     paste0("\\caption{", caption, "}"),
                     paste0("\\label{tab:", label, "}"),
                     "\\centering",
                     paste0("\\resizebox{", resize_width, "\\textwidth}{!}{%"))
    } else if(grepl("^\\\\end\\{table\\}", tex_lines[i])) {
      # Close resizebox and table
      new_lines <- c(new_lines, "}", "\\end{table}")
    } else {
      # Keep other lines as-is
      new_lines <- c(new_lines, tex_lines[i])
    }
  }
  
  # Create the final modified output
  tex_modified <- paste(new_lines, collapse = "\n")
  
  # Create directory if it doesn't exist
  file_dir <- dirname(file_name)
  if (!dir.exists(file_dir)) {
    if (verbose) {
      message("Creating directory: ", file_dir)
    }
    dir.create(file_dir, recursive = TRUE)
  }
  
  # Save the modified output
  if (verbose) {
    message("Saving table to: ", file_name)
  }
  
  cat(tex_modified, file = file_name)
  
  if (verbose) {
    message("Sensitivity table generated successfully!")
  }
  
  # Return the modified LaTeX invisibly
  invisible(tex_modified)
}
